"""
Author: Morphlng
Date: 2024-04-01 11:23:14
LastEditTime: 2025-04-21 09:15:08
LastEditors: Morphlng
Description: Carla scenario configuration helper
FilePath: /DrivingGym/src/driving_gym/environment/scenario/config/carla_scenario_config.py
"""

from __future__ import annotations

from dataclasses import asdict, dataclass

from driving_gym.environment.scenario.config.scenario_config import ScenarioConfig


class CarlaScenarioConfig(ScenarioConfig):
    def __init__(self):
        super().__init__("carla")
        self._replay_file = None
        self._replay_config = {"vehicle_as_ego": True}
        self._weather = Weather.PRESETS[0]
        self._attach_default_sensors = True
        self._enable_planner = False
        self._endless = False
        self._traffic = {
            "npc_vehicle_num": 0,
            "npc_walker_num": 0,
            "spawn_near_agents": True,
            "spawn_radius": 100,
        }

    def replay(self, replay_file: str, replay_config: dict = None):
        self._replay_file = replay_file
        self._replay_config.update(replay_config or {})
        return self

    def environment(self, map_name: str, weather=None, **kwargs):
        super().environment(map_name, **kwargs)

        def parse_weather(weather):
            if isinstance(weather, int):
                return Weather.PRESETS[weather]
            elif isinstance(weather, dict):
                return WeatherParameters(**weather)
            elif isinstance(weather, WeatherParameters):
                return weather

        if isinstance(weather, list):
            self._weather = [parse_weather(w).to_dict() for w in weather]
        elif weather is not None:
            self._weather = parse_weather(weather)
        return self

    def traffic(
        self,
        npc_vehicle_num: int | tuple[int, int] = 0,
        npc_walker_num: int | tuple[int, int] = 0,
        spawn_near_agents: bool = True,
        spawn_radius: float = 100,
    ):
        self._traffic["npc_vehicle_num"] = npc_vehicle_num
        self._traffic["npc_walker_num"] = npc_walker_num
        self._traffic["spawn_near_agents"] = spawn_near_agents
        self._traffic["spawn_radius"] = spawn_radius
        return self

    def planner(
        self, enable: bool = True, draw_path: bool = False, endless: bool = False
    ):
        """Whether to enable global planner for the agent"""
        self._enable_planner = enable
        self._draw_path = draw_path
        self._endless = endless
        return self

    def default_sensor(self, attach: bool = True):
        """Whether to attach default collision sensor and lane_invasion sensor to the agent"""
        self._attach_default_sensors = attach
        return self


@dataclass
class WeatherParameters:
    cloudiness: float = 0.0
    precipitation: float = 0.0
    precipitation_deposits: float = 0.0
    wind_intensity: float = 0.0
    sun_azimuth_angle: float = 0.0
    sun_altitude_angle: float = 0.0
    fog_density: float = 0.0
    fog_distance: float = 0.0
    wetness: float = 0.0
    fog_falloff: float = 0.0

    # The following parameters is added in Carla 0.9.12
    scattering_intensity: float = 0.0
    mie_scattering_scale: float = 0.0
    rayleigh_scattering_scale: float = 0.0331

    @classmethod
    def localize(cls, weather_parameters):
        from carla import WeatherParameters

        if not isinstance(weather_parameters, WeatherParameters):
            raise ValueError(
                "weather_parameters should be of type `carla.WeatherParameters`,"
                f"getting {type(weather_parameters)} instead"
            )

        return cls(
            cloudiness=weather_parameters.cloudiness,
            precipitation=weather_parameters.precipitation,
            precipitation_deposits=weather_parameters.precipitation_deposits,
            wind_intensity=weather_parameters.wind_intensity,
            sun_azimuth_angle=weather_parameters.sun_azimuth_angle,
            sun_altitude_angle=weather_parameters.sun_altitude_angle,
            fog_density=weather_parameters.fog_density,
            fog_distance=weather_parameters.fog_distance,
            wetness=weather_parameters.wetness,
            fog_falloff=weather_parameters.fog_falloff,
            # Compatibility with Carla 0.9.10 (Leaderboard)
            scattering_intensity=getattr(
                weather_parameters, "scattering_intensity", 0.0
            ),
            mie_scattering_scale=getattr(
                weather_parameters, "mie_scattering_scale", 0.0
            ),
            rayleigh_scattering_scale=getattr(
                weather_parameters, "rayleigh_scattering_scale", 0.0331
            ),
        )

    def to_dict(self):
        return asdict(self)


class Weather:
    """Weather presets for Simulator"""

    from carla import WeatherParameters as CarlaWeather

    PRESETS = [
        WeatherParameters.localize(CarlaWeather.ClearNoon),
        WeatherParameters.localize(CarlaWeather.CloudyNoon),
        WeatherParameters.localize(CarlaWeather.WetNoon),
        WeatherParameters.localize(CarlaWeather.WetCloudyNoon),
        WeatherParameters.localize(CarlaWeather.MidRainyNoon),
        WeatherParameters.localize(CarlaWeather.HardRainNoon),
        WeatherParameters.localize(CarlaWeather.SoftRainNoon),
        WeatherParameters.localize(CarlaWeather.ClearSunset),
        WeatherParameters.localize(CarlaWeather.CloudySunset),
        WeatherParameters.localize(CarlaWeather.WetSunset),
        WeatherParameters.localize(CarlaWeather.WetCloudySunset),
        WeatherParameters.localize(CarlaWeather.MidRainSunset),
        WeatherParameters.localize(CarlaWeather.HardRainSunset),
        WeatherParameters.localize(CarlaWeather.SoftRainSunset),
    ]
